(library (graph-algorithm)
  (export dijkstra-shortest-path
          A*
          routes->path)

  (import (except (rnrs base)
                  let-values
                  map)
          (only (guile)
                lambda* λ
                when
                ;; simple-format
                ;; with-output-to-string
                )
          (ice-9 pretty-print)
          (srfi srfi-1)  ; lists
          (srfi srfi-69)  ; hash tables
          ;; Functional Sets
          (pfds sets)
          ;; Priority Search Queues
          (pfds psqs)
          ;; Bounded Balance Trees
          (pfds bbtrees))


  (define make-empty-set
    (λ (less?)
      (make-set less?)))


  (define set-insert-multiple
    (λ (myset items)
      (cond
       [(null? items) myset]
       [else
        (set-insert-multiple (set-insert myset (car items))
                             (cdr items))])))


  (define set-empty?
    (λ (set)
      (= (set-size set) 0)))


  (define dijkstra-shortest-path
    (lambda* (start-node
              nodes
              get-neighbors
              get-neighbor-distance
              node<
              #:key
              (distance< <))
      (define init-unvisited-nodes (set-insert-multiple (make-empty-set node<) nodes))
      (define init-visited-nodes (make-empty-set node<))
      (define init-distances (alist->hash-table (map (λ (node) (cons node +inf.0)) nodes)))
      (define init-routes-table (make-hash-table eq?))
      ;; Set distance from start node to itself to 0.
      (hash-table-set! init-distances start-node 0)
      (hash-table-set! init-routes-table start-node start-node)
      ;; Visit an unvisited node with shortest known distance from start
      ;; node. Initially the start node, since all other nodes still have infinite
      ;; distance.
      (let iter ([current-node start-node]
                 [distances° init-distances]
                 [unvisited° init-unvisited-nodes]
                 [visited° init-visited-nodes]
                 [routes° init-routes-table])
        (cond
         ;; Stop, if there are no more unvisited nodes.
         [(set-empty? unvisited°)
          (values distances° routes°)]
         [else
          ;; Calculate distance to every unvisited neighbor from the start
          ;; node. The distance is the distance to current node plus distance to
          ;; the unvisted neighbor).
          (let* ([neighbors (get-neighbors current-node)]
                 ;; Only look at unvisited neighbors.
                 [unvisited-neighbors
                  (filter (λ (neighbor) (set-member? unvisited° neighbor))
                          neighbors)])
            (cond
             ;; If this particular node does not have any unvisited neighbors, go
             ;; back to visiting the next node.
             [(null? unvisited-neighbors)
              ;; Repeat until all nodes visited.
              (iter (set-fold (λ (node acc)
                                (cond [(null? acc) node]
                                      [(distance< (hash-table-ref distances° node)
                                                  (hash-table-ref distances° acc))
                                       node]
                                      [else acc]))
                              '()
                              unvisited°)
                    distances°
                    ;; Mark current node as visited.
                    (set-remove unvisited° current-node)
                    (set-insert visited° current-node)
                    routes°)]
             [else
              ;; Look at the distances to all neighbors and update distances and
              ;; routes accordingly.
              (for-each (λ (neighbor)
                          (let ([start-to-neighbor-distance
                                 (+ (hash-table-ref distances° current-node)
                                    (get-neighbor-distance current-node neighbor))])
                            ;; If distance from the start node to a neighbor node
                            ;; is less than previously known distance for that
                            ;; node, update that distance in the distances
                            ;; table. If a distance is updated, also update what
                            ;; that node's previous node on the path to the node
                            ;; is (the current node).
                            (when (distance< start-to-neighbor-distance
                                             (hash-table-ref distances° neighbor))
                              (hash-table-set! distances° neighbor start-to-neighbor-distance)
                              (hash-table-set! routes° neighbor current-node))))
                        unvisited-neighbors)
              ;; Continue with unvisted nodes.
              (iter (set-fold (λ (node acc)
                                (cond [(null? acc) node]
                                      [(distance< (hash-table-ref distances° node)
                                                  (hash-table-ref distances° acc))
                                       node]
                                      [else acc]))
                              '()
                              unvisited°)
                    distances°
                    ;; Mark current node as visited.
                    (set-remove unvisited° current-node)
                    (set-insert visited° current-node)
                    routes°)]))]))))


  (define A*
    (lambda* (start
              target
              nodes
              get-neighbors
              get-neighbor-distance
              cost-heuristic
              node<
              distance<)
      (let ([fringe (psq-set (make-psq node< <)
                             ;; Initially put the start node in the open-set,
                             ;; since we need at least some node to go on from.
                             start
                             ;; The start node has priority 0, which is the
                             ;; highest priority. It does not really matter,
                             ;; since there is only one element in the open-set.
                             0)]
            ;; routes stores the node preceding any target node on the cheapest
            ;; path to that target node. Initially only the preceding node for
            ;; the start node is set. The start node itself.
            [routes (alist->hash-table `((,start . ,start)) eq?)]
            ;; score stores the cost of the cheapest path from the start node to
            ;; other nodes as currently known. Initially it is only set for the
            ;; start node, as one does not know other costs yet.
            [cheapest-path-costs (alist->hash-table `((,start . 0)) eq?)]
            ;; Also keep track of a best estimate (calculated using the
            ;; heuristic) of the cost from the start node to the target node via
            ;; a node. Initially we have not explored any other nodes than the
            ;; start node, so the cost for any path via them to the target node
            ;; is pessimistically estimated to be infinite. Update formula:
            ;; via-node-score-estimate(node) := current-best-score(node) + heuristic(node).
            [via-node-score-estimate (make-hash-table eq?)])

        ;; Set cost estimate for start node.
        (for-each (λ (node) (hash-table-set! via-node-score-estimate node +inf.0)) nodes)
        (hash-table-set! via-node-score-estimate start (cost-heuristic start))

        (let iter ([fringe° fringe])
          ;; Initially the current-node° is the one, that is estimated to have
          ;; the lowest cost, when a path to the target contains it. Initially
          ;; this should be the start node.
          (let ([current-node° (psq-min fringe°)])
            (cond
             ;; If the current node is the target node, return the routes.
             [(eq? current-node° target) routes]
             [else
              (let ([neighbors (get-neighbors current-node°)])
                ;; Per neighbor node update the following: routes, cheapest path
                ;; cost, via node path cost estimate
                (for-each (λ (neighbor)
                            (let* ([distance (get-neighbor-distance current-node° neighbor)]
                                   ;; At first the tentative score is the distance
                                   ;; from the start to the neighbor going via the
                                   ;; current node.
                                   [tentative-score
                                    (+ (hash-table-ref cheapest-path-costs current-node°)
                                       distance)])
                              ;; If we have found a cheaper path to the neighbor
                              ;; ...
                              (when (< tentative-score (hash-table-ref cheapest-path-costs neighbor))
                                ;; ... update the preceding node on the path to
                                ;; the neighbor in the routes map ...
                                (hash-table-set! routes neighbor current-node°)
                                ;; ... and update the cheapest path costs for the
                                ;; neighbor ...
                                (hash-table-set! cheapest-path-costs neighbor tentative-score)
                                ;; ... and update the estimates of the cost of a
                                ;; path from the start node through the neighbor
                                ;; to the target node.
                                (hash-table-set! via-node-score-estimate
                                                 neighbor
                                                 (+ tentative-score
                                                    (cost-heuristic neighbor))))))
                          neighbors)
                ;; Iterate with updated fringe.
                (iter (fold (λ (node)
                              (psq-set fringe°
                                       node
                                       (hash-table-ref cheapest-path-costs node)))
                            (psq-delete fringe° current-node°)
                            neighbors)))]))))))


  ;; #;(define A*-pure
  ;;   (lambda* (start
  ;;             target
  ;;             nodes
  ;;             get-neighbors
  ;;             get-neighbor-distance
  ;;             cost-heuristic
  ;;             node<
  ;;             distance<)
  ;;     (let ([fringe (psq-set (make-psq node< <)
  ;;                            ;; Initially put the start node in the open-set,
  ;;                            ;; since we need at least some node to go on from.
  ;;                            start
  ;;                            ;; The start node has priority 0, which is the
  ;;                            ;; highest priority. It does not really matter,
  ;;                            ;; since there is only one element in the open-set.
  ;;                            0)]
  ;;           ;; routes stores the node preceding any target node on the cheapest
  ;;           ;; path to that target node. Initially only the preceding node for
  ;;           ;; the start node is set. The start node itself.
  ;;           [routes
  ;;            (bbtree-set (make-bbtree node<) start start)]
  ;;           ;; score stores the cost of the cheapest path from the start node to
  ;;           ;; other nodes as currently known. Initially it is only set for the
  ;;           ;; start node, as one does not know other costs yet.
  ;;           [cheapest-path-costs
  ;;            (bbtree-set (make-bbtree node<) start 0)]
  ;;           ;; Also keep track of a best estimate (calculated using the
  ;;           ;; heuristic) of the cost from the start node to the target node via
  ;;           ;; a node. Initially we have not explored any other nodes than the
  ;;           ;; start node, so the cost for any path via them to the target node
  ;;           ;; is pessimistically estimated to be infinite. Update formula:
  ;;           ;; via-node-score-estimate(node) := current-best-score(node) + heuristic(node).
  ;;           [via-node-score-estimate
  ;;            ;; Set cost estimate for start node.
  ;;            (psq-set
  ;;             ;; Set initial estimates for all nodes.
  ;;             (fold (λ (node queue) (psq-set queue node +inf.0))
  ;;                   (make-psq node< <)
  ;;                   nodes)
  ;;             start
  ;;             (cost-heuristic start))])
  ;;       (let iter ([current-node°
  ;;                   ;; Initially the current-node° is the one, that is estimated
  ;;                   ;; to have the lowest cost, when a path to the target
  ;;                   ;; contains it. Initially this should be the start node.
  ;;                   #;(psq-min via-node-score-estimate)
  ;;                   ;; Perhaps one can simply put start here.
  ;;                   start]
  ;;                  [fringe° fringe]
  ;;                  [routes° routes]
  ;;                  [cheapest-path-costs° cheapest-path-costs]
  ;;                  [via-node-score-estimate° via-node-score-estimate])
  ;;         (cond
  ;;          ;; If the current node is the target node, return the routes.
  ;;          [(eq? current-node° target) routes°]
  ;;          [else
  ;;           ;; TODO: do not forget to remove the current node from the fringe, since we visited it
  ;;           (let ([all-neighbors (get-neighbors current-node°)])
  ;;             ;; per neighbor node update the following:
  ;;             ;; routes, cheapest path cost, via node path cost estimate

  ;;             (iter (psq-min via-node-score-estimate)
  ;;                   ;; Add all neighbors to the fringe, if they are not yet in
  ;;                   ;; the fringe. Otherwise update their priority values.
  ;;                   (fold (λ (node queue) (psq-set queue node ... #|cheapest known path cost of node|#))
  ;;                         fringe°
  ;;                         all-neighbors)
  ;;                   #;(pretty-print (bbtree->alist (bbtree-set (bbtree-set
  ;;                   (bbtree-set (make-bbtree (lambda (k1 k2) (string<?
  ;;                   (symbol->string k1) (symbol->string k2)))) 'a 4) 'b 3)
  ;;                   'c 5)))
  ;;                   ;; --> ((a . 4) (b . 3) (c . 5))
  ;;                   )
  ;;             (let* ([distance (get-neighbor-distance current-node° neighbor)]
  ;;                    ;; At first the tentative score is the distance from the
  ;;                    ;; start to the neighbor going via the current node.
  ;;                    [tentative-score
  ;;                     (+ (hash-table-ref cheapest-path-costs° current-node°)
  ;;                        distance)])
  ;;               (cond
  ;;                ;; If we have found a cheaper path to the neighbor ...
  ;;                [(< tentative-score (hash-table-ref cheapest-path-costs° neighbor))
  ;;                 ;; if tentative_gScore < gScore[neighbor]
  ;;                 ;;     // This path to neighbor is better than any previous one. Record it!
  ;;                 ;;     cameFrom[neighbor] := current
  ;;                 ;;     gScore[neighbor] := tentative_gScore
  ;;                 ;;     fScore[neighbor] := tentative_gScore + h(neighbor)
  ;;                 ;;     if neighbor not in openSet
  ;;                 ;;         openSet.add(neighbor)

  ;;                 ;; ... update the preceding node on the path to the neighbor
  ;;                 ;; in the routes° map ...
  ;;                 (hash-table-set! routes° neighbor current-node°)
  ;;                 ;; ... and update the cheapest path costs for the neighbor
  ;;                 ;; ... (->1)
  ;;                 (hash-table-set! cheapest-path-costs° tentative-score)
  ;;                 ;; ... and update the estimates of the cost of a path from
  ;;                 ;; the start node through the neighbor to the target node.
  ;;                 #|TODO|#]
  ;;                [else ...])))
  ;;           ...])))))

  ;;     ;; // Open set is empty but goal was never reached
  ;;     ;; return failure


  (define routes->path
    (λ (routes target-node)
      "Constructs the shortest path from the start node to the
target node using the routes table."
      (let iter ([current-node° target-node] [path° (list target-node)])
        ;; (simple-format #t "current node:~a\n" current-node°)
        ;; (simple-format #t "current path:~a\n" (with-output-to-string (λ () (pretty-print path°))))
        (let ([prior-node (hash-table-ref routes current-node°)])
          (cond
           [(eq? current-node° prior-node) path°]
           [else
            (iter prior-node (cons prior-node path°))]))))))
